#!/usr/bin/env python3
"""
MCTS推理树可视化工具
简化版本 - 专注于美观的Web可视化
"""

import re
import argparse
import json
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from pathlib import Path

@dataclass
class TreeNode:
    """树节点数据结构"""
    node_id: str
    node_type: str
    q_value: float
    n_visits: int
    v_value: Optional[float]
    uct_value: float
    content: str
    depth: int
    children: Optional[List['TreeNode']] = None
    
    def __post_init__(self):
        if self.children is None:
            self.children = []

def parse_tree_file(file_path: str) -> Optional[TreeNode]:
    """解析tree文件并构建树结构"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    lines = content.split('\n')
    root = None
    node_stack = []
    
    for line in lines:
        if not line.strip():
            continue
            
        # 计算缩进深度
        depth = 0
        i = 0
        while i < len(line):
            if line[i] == '-':
                dash_count = 0
                while i < len(line) and line[i] == '-':
                    dash_count += 1
                    i += 1
                depth = dash_count // 4
                break
            elif line[i] == ' ':
                i += 1
            else:
                break
        
        stripped_line = line.lstrip(' -')
        
        # 解析节点信息
        node_match = re.match(r'\[+([^\]]+)\]+\s*(.*)', stripped_line)
        if node_match:
            node_info = node_match.group(1)
            node_content = node_match.group(2)
            
            node_params = parse_node_params(node_info)
            
            if 'id' not in node_params:
                continue
            
            node = TreeNode(
                node_id=node_params.get('id', ''),
                node_type=node_params.get('type', ''),
                q_value=node_params.get('q', 0.0),
                n_visits=node_params.get('n', 0),
                v_value=node_params.get('v'),
                uct_value=node_params.get('uct', 0.0),
                content=node_content,
                depth=depth
            )
            
            # 建立父子关系
            if depth == 0:
                if root is not None:
                    continue
                root = node
                node_stack = [node]
            else:
                while len(node_stack) > depth:
                    node_stack.pop()
                
                if node_stack:
                    parent = node_stack[-1]
                    if parent.children is None:
                        parent.children = []
                    parent.children.append(node)
                    
                    if len(node_stack) == depth:
                        node_stack.append(node)
                    else:
                        while len(node_stack) < depth:
                            node_stack.append(node_stack[-1])
                        node_stack.append(node)
    
    return root

def parse_node_params(node_info: str) -> Dict[str, Any]:
    """解析节点参数字符串"""
    params = {}
    
    # 解析节点ID和类型
    clean_info = node_info.replace('(T) ', '')
    
    id_match = re.match(r'([A-Z]+)-(\d+)', clean_info)
    if id_match:
        params['type'] = id_match.group(1)
        params['id'] = f"{id_match.group(1)}-{id_match.group(2)}"
    else:
        return params
    
    # 解析其他参数
    param_pairs = re.findall(r'([A-Z]+):\s*([^;]+)', node_info)
    for key, value in param_pairs:
        if key == 'Q':
            params['q'] = float(value)
        elif key == 'N':
            params['n'] = int(value)
        elif key == 'V':
            params['v'] = None if value == 'None' else float(value)
        elif key == 'UCT':
            # 处理特殊UCT值
            if value == '999':
                params['uct'] = 999  # 未探索节点
            elif value == '666':
                params['uct'] = 666  # 根节点
            elif value == 'inf' or value == 'Inf':
                params['uct'] = float('inf')
            else:
                try:
                    params['uct'] = float(value)
                except ValueError:
                    params['uct'] = 999  # 默认为未探索
    
    return params

def export_to_web(node: TreeNode, output_file: str, input_file: Optional[str] = None):
    """导出为交互式网页显示"""
    def node_to_dict(node: TreeNode) -> Dict[str, Any]:
        """将节点转换为字典格式"""
        uct_value = node.uct_value
        # 保持特殊UCT值不变，让前端JavaScript处理显示
        if uct_value == float('inf'):
            uct_value = 999  # 统一处理为999
        elif uct_value == float('-inf'):
            uct_value = -999
        
        return {
            'id': node.node_id,
            'type': node.node_type,
            'q_value': node.q_value,
            'n_visits': node.n_visits,
            'v_value': node.v_value,
            'uct_value': uct_value,
            'content': node.content,
            'depth': node.depth,
            'children': [node_to_dict(child) for child in (node.children or [])]
        }
    
    tree_data = node_to_dict(node)
    
    # 生成动态标题
    if input_file:
        from pathlib import Path
        input_name = Path(input_file).stem  # 获取不带扩展名的文件名
        page_title = f"{input_name} - MCTS推理树可视化"
        header_title = f"MCTS推理树可视化 - {input_name}"
    else:
        page_title = "MCTS推理树可视化"
        header_title = "MCTS推理树可视化"
    
    html_content = f"""
<!DOCTYPE html>
<html>
<head>
    <title>{page_title}</title>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <script src="https://d3js.org/d3.v7.min.js"></script>
    <style>
        * {{
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }}
        
        body {{
            font-family: 'Arial', sans-serif;
            background: #f8f9fa;
            color: #333;
            overflow: hidden;
        }}
        
        .container {{
            width: 100vw;
            height: 100vh;
            display: flex;
            flex-direction: column;
        }}
        
        .header {{
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 15px 30px;
            text-align: center;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        }}
        
        .header h1 {{
            font-size: 1.8em;
            margin: 0;
        }}
        
        .tree-container {{
            flex: 1;
            background: white;
            position: relative;
        }}
        
        .node {{
            cursor: pointer;
        }}
        
        .node circle {{
            fill: #fff;
            stroke-width: 2px;
            r: 20;
        }}
        
        .node.U circle {{ fill: #e3f2fd; stroke: #2196f3; }}
        .node.TS circle {{ fill: #e8f5e8; stroke: #4caf50; }}
        .node.DA circle {{ fill: #fff3e0; stroke: #ff5722; }}
        .node.SQ circle {{ fill: #f5f5f5; stroke: #757575; }}
        .node.RS circle {{ fill: #fce4ec; stroke: #e91e63; }}
        .node.RU circle {{ fill: #f3e5f5; stroke: #9c27b0; }}
        
        /* UCT状态样式 */
        .node.unexplored circle {{
            stroke-dasharray: 5,5;  /* 虚线边框表示未探索 */
            opacity: 0.7;
        }}
        
        .node.root circle {{
            stroke-width: 4px;  /* 加粗边框表示根节点 */
        }}
        
        .node:hover circle {{
            stroke-width: 3px;
            filter: drop-shadow(0 2px 4px rgba(0,0,0,0.2));
        }}
        
        .node.root:hover circle {{
            stroke-width: 5px;  /* 根节点悬停时更粗 */
        }}
        
        .link {{
            fill: none;
            stroke: #9ecae1;
            stroke-width: 2px;
            opacity: 0.7;
        }}
        
        .node-text {{
            font-size: 11px;
            text-anchor: middle;
            pointer-events: none;
            font-weight: bold;
            fill: #333;
        }}
        
        .tooltip {{
            position: absolute;
            background: rgba(0, 0, 0, 0.9);
            color: white;
            padding: 8px 10px;
            border-radius: 6px;
            font-size: 12px;
            max-width: 600px;
            max-height: 400px;
            overflow-y: auto;
            z-index: 1000;
            display: none;
            box-shadow: 0 4px 20px rgba(0,0,0,0.3);
            line-height: 1.2;
            white-space: pre-wrap;
        }}
        
        .tooltip-title {{
            font-weight: bold;
            font-size: 14px;
            margin-bottom: 4px;
            color: #4fc3f7;
        }}
        
        .tooltip-content {{
            margin-bottom: 1px;
            line-height: 1.2;
        }}
        
        .tooltip-content strong {{
            color: #81c784;
            margin-right: 2px;
        }}
        
        .tooltip-stats {{
            display: flex;
            gap: 12px;
            margin-bottom: 4px;
            flex-wrap: wrap;
        }}
        
        .tooltip-stat {{
            font-size: 11px;
            line-height: 1.1;
        }}
        
        .tooltip-content-text {{
            margin-top: 4px;
            padding-left: 4px;
            font-size: 11px;
            line-height: 1.2;
            border-left: 2px solid #4fc3f7;
            padding-left: 6px;
        }}
        
        .legend {{
            position: absolute;
            top: 20px;
            right: 20px;
            background: rgba(255,255,255,0.95);
            padding: 15px;
            border-radius: 8px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            font-size: 12px;
            max-width: 250px;
        }}
        
        .legend-title {{
            font-weight: bold;
            margin-bottom: 10px;
            color: #333;
        }}
        
        .legend-item {{
            display: flex;
            align-items: center;
            margin-bottom: 8px;
        }}
        
        .legend-color {{
            width: 16px;
            height: 16px;
            border-radius: 50%;
            margin-right: 8px;
            border: 2px solid;
        }}
        
        .legend-color.U {{ background: #e3f2fd; border-color: #2196f3; }}
        .legend-color.TS {{ background: #e8f5e8; border-color: #4caf50; }}
        .legend-color.DA {{ background: #fff3e0; border-color: #ff5722; }}
        .legend-color.SQ {{ background: #f5f5f5; border-color: #757575; }}
        .legend-color.RS {{ background: #fce4ec; border-color: #e91e63; }}
        .legend-color.RU {{ background: #f3e5f5; border-color: #9c27b0; }}
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>{header_title}</h1>
        </div>
        
        <div class="tree-container">
            <div class="legend">
                <div class="legend-title">节点类型</div>
                <div class="legend-item">
                    <div class="legend-color U"></div>
                    <span>User (用户)</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color TS"></div>
                    <span>Thought Step (思考步骤) - A1</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color DA"></div>
                    <span>Direct Answer (直接回答) - A2</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color SQ"></div>
                    <span>Sub Question (子问题) - A3</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color RS"></div>
                    <span>Re-Answer (重新回答) - A4</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color RU"></div>
                    <span>Rephrase User (重述用户) - A5</span>
                </div>
                
                <div class="legend-title" style="margin-top: 15px;">UCT状态</div>
                <div style="font-size: 11px; margin-bottom: 8px;">
                    <div>• 虚线边框：未探索节点 (UCT=999)</div>
                    <div>• 加粗边框：根节点 (UCT=666)</div>
                    <div>• 普通边框：已探索节点 (UCT=数值)</div>
                </div>
            </div>
            
            <div class="tooltip" id="tooltip"></div>
            <svg id="tree-svg"></svg>
        </div>
    </div>

    <script>
        const treeData = {json.dumps(tree_data, ensure_ascii=False, indent=2)};
        
        // 设置SVG尺寸
        const container = document.querySelector('.tree-container');
        const width = container.clientWidth;
        const height = container.clientHeight;
        
        const svg = d3.select('#tree-svg')
            .attr('width', width)
            .attr('height', height);
        
        const g = svg.append('g');
        
        // 创建树布局
        const tree = d3.tree()
            .size([width - 100, height - 100])
            .separation((a, b) => {{
                return a.parent === b.parent ? 1.5 : 2;
            }});
        
        // 转换数据
        const root = d3.hierarchy(treeData);
        tree(root);
        
        // 调整位置
        root.each(d => {{
            d.x += 50;
            d.y += 50;
        }});
        
        // 绘制连线
        g.selectAll('.link')
            .data(root.links())
            .enter()
            .append('path')
            .attr('class', 'link')
            .attr('d', d3.linkVertical()
                .x(d => d.x)
                .y(d => d.y)
            );
        
        // 绘制节点
        const nodes = g.selectAll('.node')
            .data(root.descendants())
            .enter()
            .append('g')
            .attr('class', d => {{
                let classes = `node ${{d.data.type}}`;
                // 根据UCT值添加状态类
                if (d.data.uct_value === 999) {{
                    classes += ' unexplored';
                }} else if (d.data.uct_value === 666) {{
                    classes += ' root';
                }}
                return classes;
            }})
            .attr('transform', d => `translate(${{d.x}},${{d.y}})`);
        
        nodes.append('circle');
        
        nodes.append('text')
            .attr('class', 'node-text')
            .attr('dy', 4)
            .text(d => d.data.id);
        
        // 添加工具提示
        const tooltip = d3.select('#tooltip');
        let tooltipTimeout;
        
        // 鼠标悬停显示完整信息
        nodes.on('mouseover', function(event, d) {{
            clearTimeout(tooltipTimeout);
            
            // 格式化UCT值显示
            let uctDisplay;
            if (d.data.uct_value === 999) {{
                uctDisplay = 'Unexplored';  // 未探索节点
            }} else if (d.data.uct_value === 666) {{
                uctDisplay = 'Root';        // 根节点
            }} else if (d.data.uct_value === -999) {{
                uctDisplay = '-Inf';
            }} else {{
                uctDisplay = d.data.uct_value.toFixed(2);
            }}
            
            const content = `
                <div class="tooltip-title">${{d.data.id}} (${{d.data.type}})</div>
                <div class="tooltip-stats">
                    <div class="tooltip-stat"><strong>Q:</strong>${{d.data.q_value.toFixed(2)}}</div>
                    <div class="tooltip-stat"><strong>N:</strong>${{d.data.n_visits}}</div>
                    <div class="tooltip-stat"><strong>V:</strong>${{d.data.v_value !== null ? d.data.v_value.toFixed(2) : 'None'}}</div>
                    <div class="tooltip-stat"><strong>UCT:</strong>${{uctDisplay}}</div>
                </div>
                <div class="tooltip-content-text">${{d.data.content}}</div>
            `;
            
            tooltip.html(content)
                .style('display', 'block')
                .style('left', Math.min(event.pageX + 10, window.innerWidth - 620) + 'px')
                .style('top', Math.min(event.pageY - 10, window.innerHeight - 420) + 'px');
        }})
        .on('mouseout', function() {{
            tooltipTimeout = setTimeout(() => {{
                if (tooltip.style('display') === 'block') {{
                    // 检查鼠标是否在tooltip上
                    const tooltipElement = document.getElementById('tooltip');
                    if (!tooltipElement.matches(':hover')) {{
                        tooltip.style('display', 'none');
                    }}
                }}
            }}, 300);
        }});
        

        
        // 鼠标进入tooltip时保持显示
        tooltip.on('mouseenter', function() {{
            clearTimeout(tooltipTimeout);
        }})
        .on('mouseleave', function() {{
            tooltip.style('display', 'none');
        }});
        
        // 添加缩放和拖拽
        const zoom = d3.zoom()
            .scaleExtent([0.1, 3])
            .on('zoom', function(event) {{
                g.attr('transform', event.transform);
            }});
        
        svg.call(zoom);
        
        // 适应窗口大小变化
        window.addEventListener('resize', function() {{
            const newWidth = container.clientWidth;
            const newHeight = container.clientHeight;
            svg.attr('width', newWidth).attr('height', newHeight);
            
            tree.size([newWidth - 100, newHeight - 100]);
            tree(root);
            
            root.each(d => {{
                d.x += 50;
                d.y += 50;
            }});
            
            g.selectAll('.link')
                .attr('d', d3.linkVertical()
                    .x(d => d.x)
                    .y(d => d.y)
                );
            
            g.selectAll('.node')
                .attr('transform', d => `translate(${{d.x}},${{d.y}})`);
        }});
    </script>
</body>
</html>
    """
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(html_content)
    
    print(f"网页可视化已保存到: {output_file}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='MCTS推理树可视化工具')
    parser.add_argument('input_file', help='输入的.tree文件路径')
    parser.add_argument('-o', '--output', default=None, 
                       help='输出HTML文件路径 (默认: 与输入文件同名但扩展名为.html)')
    
    args = parser.parse_args()
    
    # 检查输入文件
    input_path = Path(args.input_file)
    if not input_path.exists():
        print(f"错误: 输入文件 {args.input_file} 不存在")
        return
    
    # 如果没有指定输出文件，则使用与输入文件同名但扩展名为.html的文件
    if args.output is None:
        args.output = str(input_path.with_suffix('.html'))
    
    print(f"正在解析文件: {args.input_file}")
    
    # 解析树文件
    root = parse_tree_file(args.input_file)
    if not root:
        print("错误: 无法解析树文件")
        return
    
    print(f"解析成功，根节点: {root.node_id}")
    
    # 生成可视化
    export_to_web(root, args.output, args.input_file)

if __name__ == "__main__":
    main() 